

################################################################################################
################################################################################################


import numpy
import math

import os
import sys

sys.path.insert(0, os.path.dirname(__file__))
import pyglet
from pyglet import gl
from shader import FragmentShader, ShaderError, ShaderProgram, VertexShader

from ..elements import elements


################################################################################################
################################################################################################


ROTATE, TRANSLATE, SPIN = range(3)


################################################################################################
################################################################################################


def vec(*args):
    return (gl.GLfloat * len(args))(*args)


################################################################################################
################################################################################################


class Beam(pyglet.window.Window):


################################################################################################


    def initialize(self):

        gl.glEnable(gl.GL_DEPTH_TEST)
        gl.glDepthFunc(gl.GL_LEQUAL)
        gl.glEnable(gl.GL_ALPHA_TEST)
        gl.glAlphaFunc(gl.GL_EQUAL, 1.0)   
        
        try: 

            fshader = FragmentShader([open(os.path.join(os.path.dirname(__file__), "fragment.glsl"), 'r').read()])
            vshader = VertexShader([open(os.path.join(os.path.dirname(__file__), "vertex.glsl"), 'r').read()])
            self.shader = ShaderProgram(fshader, vshader)
            self.shader.use()

            sphereSampler = gl.glGetUniformLocation(self.shader.id, "sphere");        
            depthSampler = gl.glGetUniformLocation(self.shader.id, "depth");        

            gl.glUniform1i(sphereSampler, 0)
            gl.glUniform1i(depthSampler, 1)

            spherePNG = pyglet.image.load(os.path.join(os.path.dirname(__file__), "sphere1.png"))
            self.sphereTexture = spherePNG.get_texture()
            depthPNG = pyglet.image.load(os.path.join(os.path.dirname(__file__), "depth.png"))
            self.depthTexture = depthPNG.get_texture()

            gl.glActiveTexture(gl.GL_TEXTURE0)
            gl.glEnable(self.sphereTexture.target)
            gl.glBindTexture(self.sphereTexture.target, self.sphereTexture.id)
            gl.glActiveTexture(gl.GL_TEXTURE1)
            gl.glEnable(self.depthTexture.target)
            gl.glBindTexture(self.depthTexture.target, self.depthTexture.id)
            
        except:

            spherePNG = pyglet.image.load(os.path.join(os.path.dirname(__file__), "sphere1.png"))
            self.sphereTexture = spherePNG.get_texture()

            gl.glActiveTexture(gl.GL_TEXTURE0)
            gl.glEnable(self.sphereTexture.target)
            gl.glBindTexture(self.sphereTexture.target, self.sphereTexture.id)

            print "Could not compile depth sprite shaders."





################################################################################################


    def __init__(self, *args, **kwargs):
        pyglet.window.Window.__init__(self, resizable = True, fullscreen = False, vsync = True, width = 512, height = 512)
        self.scale = 1.0
        self.mouseActivity = None
        self.zoom = 10.0
        self.keys = pyglet.window.key.KeyStateHandler()
        self.push_handlers(self.keys)
        self.initialize()
        self.background = 1.0
        self.rotation = numpy.identity((3))
        self.targetRotation = numpy.identity((3))
        self.translation = numpy.array([0.0, 0.0, 0.0])
        self.middle = None


################################################################################################


    def on_mouse_press(self, x, y, button, modifiers):
        if button & pyglet.window.mouse.LEFT:
            self.mouseActivity = ROTATE
        elif button & pyglet.window.mouse.RIGHT:
            self.mouseActivity = TRANSLATE
        elif button & pyglet.window.mouse.MIDDLE:
            self.mouseActivity = SPIN


################################################################################################


    def on_mouse_drag(self, x, y, dx, dy, buttons, modifiers):
        if self.mouseActivity == ROTATE:
            self.rotX(-dy * 0.0078125)
            self.rotY(dx * 0.0078125)
        elif self.mouseActivity == SPIN:
            self.rotZ(-dx * 0.0078125)
        elif self.mouseActivity == TRANSLATE:
            self.translation[0] += (float(dx) / self.width) * (2 * self.zoom * (float(self.width)/self.height))
            self.translation[1] += dy * (self.zoom / self.height) * 2.0
        self.invalid = True


################################################################################################


    def on_mouse_release(self, x, y, button, modifiers):
        if self.mouseActivity == ROTATE:
            self.mouseActivity = None


################################################################################################


    def on_mouse_scroll(self, x, y, scroll_x, scroll_y):
        if self.keys[pyglet.window.key.R]:
            if scroll_y < 0:
                self.scale *= 0.9
            else:
                self.scale *= 1.1
        else:
            if scroll_y > 0:
                self.zoom *= 0.9
            else:
                self.zoom *= 1.1
        self.invalid = True


################################################################################################


    def clear(self):
        gl.glClearColor(self.background, self.background, self.background, 1.0) 
        gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)


################################################################################################


    def on_draw(self):

        # If we don't have a speck, get out of here.
        try:
            self.speck
        except:
            return

        # Update the rotation.  Used for smooth rotations.
        self.rotation += (self.targetRotation - self.rotation) * 1.0 
        
        # Handy variable for the length.
        l = len(self.speck)
        
        # Get a copy of r.
        rd = self.speck.r.copy()
        
        # If we have not calculated the middle of the speck, calculate it.
        if self.middle == None:
            minx = min(rd[:, 0])                         
            miny = min(rd[:, 1])                         
            minz = min(rd[:, 2])
            maxx = max(rd[:, 0])
            maxy = max(rd[:, 1])
            maxz = max(rd[:, 2])
            midx = minx + (maxx - minx) * 0.5
            midy = miny + (maxy - miny) * 0.5
            midz = minz + (maxz - minz) * 0.5
            self.middle = numpy.array([midx, midy, midz])
            
        # Center the middle of the speck.            
        rd -= self.middle
        
        # Perform the rotation (matrix multiply according to...)
        # http://thread.gmane.org/gmane.comp.python.numeric.general/20360/focus=21033
        rd = numpy.sum(self.rotation[...,:,:]*rd[...,numpy.newaxis,:], axis=-1) 
        
        # Translate the speck.
        rd += self.translation
        
        # Sort the atoms in z.
        indeces = range(l)
        indeces = sorted(indeces, key = lambda i: rd[i][2], reverse = True)
        
        # Create the quad vertices.
        r = []
        for i in indeces:
            ss = elements[self.speck.symbols[i]]['radius'] * self.scale
            r.extend([-ss + rd[i][0], -ss + rd[i][1], rd[i][2], 
                      -ss + rd[i][0],  ss + rd[i][1], rd[i][2], 
                       ss + rd[i][0],  ss + rd[i][1], rd[i][2], 
                       ss + rd[i][0], -ss + rd[i][1], rd[i][2]])
        
        # Create the texture coordinates for the quads.
        uv = [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0] * l
        
        # Determine the radii for the depth shader.
        nearFar = 1000.0
        invNearFarScale = self.scale / (2.0 * nearFar)
        radii = []
        for i in indeces:
            radii.extend([elements[self.speck.symbols[i]]['radius'] * invNearFarScale, 0.0, 0.0] * 4)
            
        # Create the color data.
        colors = []
        for i in indeces:
            colors.extend([elements[self.speck.symbols[i]]['color'][0], 
                           elements[self.speck.symbols[i]]['color'][1], 
                           elements[self.speck.symbols[i]]['color'][2]] * 4)

        # Fill the vertex list.
        vl = pyglet.graphics.vertex_list(l * 4, ('v3f', r), ('t2f', uv), ('c3f', colors), ('n3f', radii))

        # Clear the screen and set up the projection.
        self.clear()
        gl.glViewport(0, 0, self.width, self.height)
        gl.glMatrixMode(gl.GL_PROJECTION)
        gl.glLoadIdentity()
        zwhscale = self.zoom * (float(self.width)/self.height)
        gl.glOrtho(-zwhscale, zwhscale, -self.zoom, self.zoom, -nearFar, nearFar)
        gl.glMatrixMode(gl.GL_MODELVIEW)
        gl.glLoadIdentity()
        
        # Draw the depth sprites.
        vl.draw(gl.GL_QUADS)
        
        # Nix the vertex list.
        vl.delete()
        
        # Set the invalid flag to false so that on_draw is not called in the pyglet idle method.
        self.invalid = False


################################################################################################


    def rotX(self, theta):
        ct = math.cos(theta)
        st = math.sin(theta)
        m = numpy.array([[1, 0, 0], [0, ct, -st], [0, st, ct]])
        self.targetRotation = numpy.dot(m, self.targetRotation)
    def rotY(self, theta):
        ct = math.cos(theta)
        st = math.sin(theta)
        m = numpy.array([[ct, 0, st], [0, 1, 0], [-st, 0, ct]])
        self.targetRotation = numpy.dot(m, self.targetRotation)
    def rotZ(self, theta):
        ct = math.cos(theta)
        st = math.sin(theta)
        m = numpy.array([[ct, -st, 0], [st, ct, 0], [0, 0, 1]])
        self.targetRotation = numpy.dot(m, self.targetRotation)
        

























